# ai_transaction_detector.py
import numpy as np
import pandas as pd
from sklearn.ensemble import IsolationForest
from datetime import datetime
import joblib  # 用于保存模型

class AITransactionDetector:
    def __init__(self, contamination=0.1):
        self.model = IsolationForest(contamination=contamination, random_state=42)
        self.is_fitted = False

    def prepare_features(self, df):
        """提取特征"""
        X = df[['duration_sec', 'wait_time_ms', 'rows_affected', 'lock_count']].copy()
        # 标准化（可选）
        X = (X - X.mean()) / X.std()
        return X.fillna(0)

    def train(self, log_data):
        """
        训练模型
        log_data: list of dict or DataFrame
        """
        if isinstance(log_data, list):
            df = pd.DataFrame(log_data)
        else:
            df = log_data.copy()

        X = self.prepare_features(df)
        self.model.fit(X)
        self.is_fitted = True
        print("✅ AI 模型训练完成")

        # 保存模型和数据
        joblib.dump(self.model, "transaction_model.pkl")
        df.to_csv("training_logs.csv", index=False)
        return self

    def predict(self, new_tx):
        """
        预测单个事务是否为异常长事务
        new_tx: dict like {'duration_sec': 300, 'wait_time_ms': 500, ...}
        """
        if not self.is_fitted:
            raise Exception("模型未训练，请先调用 train()")

        df = pd.DataFrame([new_tx])
        X = self.prepare_features(df)
        pred = self.model.predict(X)[0]  # 1 正常，-1 异常
        score = self.model.decision_function(X)[0]

        is_anomaly = pred == -1
        confidence = abs(score)

        return {
            "is_long_running": is_anomaly,
            "confidence": float(confidence),
            "risk_level": "高" if is_anomaly else "低",
            "suggestion": "建议监控或考虑回滚" if is_anomaly else "运行正常"
        }

# 示例使用
if __name__ == "__main__":
    # 模拟训练数据
    logs = [
        {"tx_id": "tx_001", "duration_sec": 10, "wait_time_ms": 5, "rows_affected": 100, "lock_count": 1},
        {"tx_id": "tx_002", "duration_sec": 5, "wait_time_ms": 2, "rows_affected": 50, "lock_count": 1},
        {"tx_id": "tx_003", "duration_sec": 600, "wait_time_ms": 800, "rows_affected": 10, "lock_count": 5},  # 异常
        {"tx_id": "tx_004", "duration_sec": 800, "wait_time_ms": 1000, "rows_affected": 5, "lock_count": 8},  # 异常
        {"tx_id": "tx_005", "duration_sec": 8, "wait_time_ms": 3, "rows_affected": 200, "lock_count": 2},
    ]

    detector = AITransactionDetector(contamination=0.3)
    detector.train(logs)

    # 测试新事务
    new_tx = {
        "duration_sec": 400,
        "wait_time_ms": 600,
        "rows_affected": 20,
        "lock_count": 6
    }
    result = detector.predict(new_tx)
    print("🔍 检测结果:", result)